package cn.da.shuai.cool.ai.controller;

import cn.da.shuai.cool.ai.model.dto.ChatDTO;
import cn.da.shuai.cool.ai.tool.WeatherTool;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatModel;
import lombok.Setter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.SafeGuardAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.document.Document;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.List;

@RestController
@RequestMapping("/chat/stream")
public class ChatStreamController {

    private final DashScopeChatModel dashScopeChatModel;

    private final ChatClient chatClient;

    @Setter(onMethod_ = @Autowired)
    private Advisor messageChatMemoryAdvisor;

    @Setter(onMethod_ = @Autowired)
    private WeatherTool weatherTool;

    @Setter(onMethod_ = @Autowired)
    private VectorStore vectorStore;

    /**
     * 通过构造函数注入的方式注入不同大模型的实例
     *
     * @param dashScopeChatModel 阿里模型实例
     * @param ollamaChatModel    ollama模型实例
     */
    public ChatStreamController(DashScopeChatModel dashScopeChatModel, OllamaChatModel ollamaChatModel) {
        this.chatClient = ChatClient.builder(ollamaChatModel).build();
        this.dashScopeChatModel = dashScopeChatModel;
    }

    @GetMapping(value = "/query", produces = "text/html;charset=utf-8")
    public Flux<String> chatModel(@RequestParam(defaultValue = "你是谁") String question) {
        return dashScopeChatModel.stream(question);
    }

    /**
     * 流式聊天接口，返回SSE(Server-Sent Events)格式的数据流，
     * 如果想返回网页数据，请将 produces 参数设置为 "text/html;charset=utf-8"
     *
     * @param question 问题
     * @return 聊天内容
     */
    @GetMapping(value = "/query-stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> chatStream(@RequestParam(defaultValue = "你是谁") String question) {
        // 使用chatClient构建prompt，流式返回聊天内容
        return chatClient.prompt("我们名字是小帅")
                // 设置系统提示
                .system(system -> system
                        .text("你是乐于助人的人工助手。 当前时间：{now}")
                        .param("now", LocalDateTime.now())
                )
                .user(question)  // 设置用户问题
                .advisors(new SimpleLoggerAdvisor()) // 添加日志拦截器, 需要在配置文件中进行配置日志路径
                .stream()        // 开启流式处理
                .content()      // 获取内容流
                .doOnNext(System.out::println)  // 打印每条响应内容
                .concatWith(Flux.just("[END]"));  // 在流结束时添加结束标记
    }

    /**
     * 处理聊天记忆查询请求
     *
     * @param dto 包含聊天问题和会话ID的请求数据传输对象
     * @return 返回包含AI回复内容的字符串流
     */
    @PostMapping(value = "/query", produces = "text/html;charset=utf-8")
    public Flux<String> chatMemory(@RequestBody ChatDTO dto) {
        Logger logger = LoggerFactory.getLogger(this.getClass());

        List<String> sensitiveWords = getSensitiveWords(); // 从配置或服务中获取敏感词列表

        String now = LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss"));

        return chatClient
                .prompt()
                // 设置系统提示词，包含当前时间参数
                .system(system -> system.text("你是乐于助人人工智能助手。 当前时间：{now}")
                        .param("now", now))
                .user(dto.getQuestion())
                // 配置聊天记忆顾问，用于维护会话上下文
                .advisors(advisor -> advisor
                        .advisors(messageChatMemoryAdvisor)
                        .param(ChatMemory.CONVERSATION_ID, dto.getConversationId())
                )
                // 添加日志记录顾问
                .advisors(new SimpleLoggerAdvisor())
                // 添加敏感词过滤顾问
                .advisors(new SafeGuardAdvisor(
                        sensitiveWords,
                        "敏感词提示，请勿输入敏感词",
                        Advisor.DEFAULT_CHAT_MEMORY_PRECEDENCE_ORDER + 1)) // 确保顺序合理
                .advisors(
                        QuestionAnswerAdvisor
                                .builder(vectorStore)
                                .searchRequest(SearchRequest.builder()
                                        .query(dto.getQuestion())
                                        .topK(3)
                                        .similarityThreshold(0.5f)
                                        .build())
                                .build()
                )
                .tools(weatherTool)
                .stream()
                .content()
                // 实时打印AI回复内容到日志
                .doOnNext(content -> logger.info("AI response chunk: {}", content))
                // 捕获并记录异常
                .doOnError(error -> logger.error("Error occurred during chat streaming", error));
    }

    @GetMapping("/search")
    public List<Document> search(@RequestParam("query") String query) {
        return vectorStore.similaritySearch(
                SearchRequest
                        .builder()
                        .topK(3)
                        .query(query)
                        .similarityThreshold(0.5f)
                        .build()
        );
    }

    // 示例：从配置或服务中获取敏感词列表
    private List<String> getSensitiveWords() {
        // 这里可以从 application.yml 加载，或者调用远程服务获取
        return List.of("傻"); // 实际应替换为动态加载逻辑
    }


}
